""" PointNet classification infer script. """

import argparse
import os
import sys
import importlib
import mindspore
import mindspore.dataset as ds
import mindspore.ops as ops


from mindspore import context, load_checkpoint, load_param_into_net
from mind3d.dataset.ModelNet40 import ModelNet40Dataset
from mindspore.communication.management import init
from mindspore.common import set_seed
from mindspore.context import ParallelMode
from mind3d.utils.load_yaml import load_yaml
from visualization.visualization import show_pointcloud_cls
from mind3d.models.dgcnn import DGCNN_cls
from mind3d.models.pointnet import PointNet_cls
from mind3d.models.pointnet2 import Pointnet2clsModelSSG
from mind3d.models.PointTransformer import PointTransformerCls

set_seed(1)
BASE_DIR = os.path.dirname(os.path.abspath(__file__))  # '/ms3d/example/pointnet
sys.path.append(os.path.dirname(os.path.dirname(BASE_DIR)))
ROOT_DIR = os.path.join(os.path.dirname(os.path.dirname(BASE_DIR)), 'models')  # '/ms3d/models
sys.path.append(ROOT_DIR)

def infer(opt):
    """infer."""
    # device.
    device_id = int(os.getenv('DEVICE_ID', '0'))
    device_num = int(os.getenv('RANK_SIZE', '0'))

    if not opt['device_target'] in ("Ascend", "GPU"):
        raise ValueError("Unsupported platform {}".format(opt['device_target']))

    if opt['device_target'] == "Ascend":
        context.set_context(mode=context.GRAPH_MODE,
                            device_target="Ascend",
                            save_graphs=False,
                            device_id=device_id)
        context.set_context(max_call_depth=2048)

    else:
        context.set_context(mode=context.GRAPH_MODE,
                            device_target="GPU",
                            save_graphs=False,
                            device_id=device_id,
                            max_call_depth=2048)

    # run distribute.
    if opt['run_distribute']:
        if opt['device_target'] == "Ascend":
            if device_num > 1:
                init()
                context.set_auto_parallel_context(
                    parallel_mode=ParallelMode.DATA_PARALLEL, gradients_mean=True)
        else:
            if device_num > 1:
                mindspore.dataset.config.set_enable_shared_mem(False)
                context.set_auto_parallel_context(
                    parallel_mode=context.ParallelMode.DATA_PARALLEL,
                    gradients_mean=True,
                    device_num=device_num)
                mindspore.common.set_seed(1234)
                init()
            else:
                context.set_context(device_id=device_id)

    # Data pipeline.
    dataset = ModelNet40Dataset(root_path=opt['datasets']['val'].get('data_url'),
                                split=opt['datasets']['val'].get('split'),
                                num_points=opt['datasets']['val'].get('resize'),
                                use_norm=opt['datasets']['val'].get('use_norm'))

    dataset_val = ds.GeneratorDataset(dataset, ["data", "label"], shuffle=False)
    dataset_val = dataset_val.batch(1, drop_remainder=True)

    '''MODEL LOADING'''
    #create model
    network_name = opt["model_name"]

    if network_name == "DGCNN_cls":
        network = DGCNN_cls(opt, output_channels=opt['train'].get("num_classes"))
    elif network_name == "PointNet_cls":
        network = PointNet_cls(k=opt['train'].get('num_classes'))
    elif network_name == "Pointnet2clsModelSSG":
        network = Pointnet2clsModelSSG(normal_channel=opt['datasets']['train'].get('use_norm'))
    elif network_name == "PointTransformerCls":
        network = PointTransformerCls()

    # Load checkpoint file for ST test.
    param_dict = load_checkpoint(opt['val']['ckpt_file'])
    load_param_into_net(network, param_dict)

    # Begin to infer.
    for data in dataset_val.create_dict_iterator():

        points, target = data['data'], data['label']
        _, _, C = points.shape #B N 3

        cls_pred = network(points) # B 40

        pred = ops.ExpandDims()(cls_pred,    0)
        preds = pred.argmax(2)

        points = points.view(-1, C).asnumpy()
        cls_pred = preds.squeeze(0).view(-1).asnumpy()

        show_pointcloud_cls(points, label=cls_pred, save=None)


if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='Infer.')
    parser.add_argument('-opt', type=str, default='/home/cxh/下载/fin/mind3d/configs/point_transformer/PointTransformerCls.yaml', help='Path to option YAML file.')
    args = parser.parse_known_args()[0]
    opt = load_yaml(args.opt)
    infer(opt)
